package com.zz.common.utils;

import com.zz.common.exception.TunnelException;
import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.Supplier;

@Slf4j
public class LockUtil {
    private static final ConcurrentHashMap<String, ReentrantLock> lockMap = new ConcurrentHashMap<>();

    /**
     * 获取与给定键关联的锁，如果不存在则创建一个新的锁，默认为非公平锁。
     *
     * @param key 锁的唯一标识符
     * @return 与键关联的ReentrantLock实例
     */
    public static ReentrantLock getLock(String key) {
        return getLock(key, false);
    }

    /**
     * 获取与给定键关联的锁，如果不存在则创建一个新的锁，并可以指定是否为公平锁。
     *
     * @param key  锁的唯一标识符
     * @param fair 是否为公平锁
     * @return 与键关联的ReentrantLock实例
     */
    public static ReentrantLock getLock(String key, boolean fair) {
        return lockMap.computeIfAbsent(key, k -> new ReentrantLock(fair));
    }

    /**
     * 尝试获取锁并执行给定的任务。此方法会自动处理锁的获取和释放。
     *
     * @param key  锁的唯一标识符
     * @param task 需要在持有锁的情况下执行的任务
     * @param <T>  任务返回的结果类型
     * @return 任务执行结果
     * @throws Exception 如果任务执行过程中抛出异常，则传递该异常
     */
    public static <T> T executeWithLock(String key, Supplier<T> task) throws Exception {
        return executeWithLock(key, task, false, 0L, TimeUnit.MILLISECONDS);
    }

    /**
     * 尝试在指定时间内获取锁并执行给定的任务。此方法会自动处理锁的获取和释放。
     *
     * @param key     锁的唯一标识符
     * @param task    需要在持有锁的情况下执行的任务
     * @param timeout 等待锁的最大时间
     * @param unit    时间单位
     * @param <T>     任务返回的结果类型
     * @return 任务执行结果
     * @throws Exception 如果任务执行过程中抛出异常，则传递该异常
     */
    public static <T> T executeWithLock(String key, Supplier<T> task, long timeout, TimeUnit unit) throws Exception {
        return executeWithLock(key, task, false, timeout, unit);
    }

    /**
     * 尝试获取锁并执行给定的任务。此方法会自动处理锁的获取和释放，并可以选择是否为公平锁以及等待获取锁的时间。
     *
     * @param key     锁的唯一标识符
     * @param task    需要在持有锁的情况下执行的任务
     * @param fair    是否为公平锁
     * @param timeout 等待锁的最大时间
     * @param unit    时间单位
     * @param <T>     任务返回的结果类型
     * @return 任务执行结果
     * @throws Exception 如果任务执行过程中抛出异常，则传递该异常
     */
    public static <T> T executeWithLock(String key, Supplier<T> task, boolean fair, long timeout, TimeUnit unit) throws Exception {
        ReentrantLock lock = getLock(key, fair);
        try {
            if (timeout > 0) {
                if (!lock.tryLock(timeout, unit)) {
                    throw new TunnelException("Could not acquire lock within the specified time.");
                }
            } else {
                lock.lock();
            }
            log.info(Thread.currentThread().getName() + " acquired lock for " + key);
            return task.get(); // 执行任务
        } finally {
            lock.unlock();
            log.info(Thread.currentThread().getName() + " released lock for " + key);
        }
    }

    /**
     * 尝试获取锁。如果不能立即获取，则返回false。
     *
     * @param key 锁的唯一标识符
     * @return 是否成功获取锁
     */
    public static boolean tryLock(String key) {
        ReentrantLock lock = getLock(key);
        return lock.tryLock();
    }

    /**
     * 尝试在指定时间内获取锁。如果超时未获取到，则返回false。
     *
     * @param key     锁的唯一标识符
     * @param timeout 等待锁的最大时间
     * @param unit    时间单位
     * @return 是否成功获取锁
     * @throws InterruptedException 如果线程被中断
     */
    public static boolean tryLock(String key, long timeout, TimeUnit unit) throws InterruptedException {
        ReentrantLock lock = getLock(key);
        return lock.tryLock(timeout, unit);
    }

    /**
     * 检查锁是否被当前线程持有。
     *
     * @param key 锁的唯一标识符
     * @return 当前线程是否持有该锁
     */
    public static boolean isHeldByCurrentThread(String key) {
        ReentrantLock lock = lockMap.get(key);
        return lock != null && lock.isHeldByCurrentThread();
    }

    /**
     * 检查是否有其他线程正在等待锁。
     *
     * @param key 锁的唯一标识符
     * @return 是否有其他线程等待该锁
     */
    public static int getQueueLength(String key) {
        ReentrantLock lock = lockMap.get(key);
        return lock != null ? lock.getQueueLength() : 0;
    }

    /**
     * 清除指定键对应的锁，通常用于清理不再使用的锁。
     *
     * @param key 锁的唯一标识符
     */
    public static void clearLock(String key) {
        ReentrantLock removedLock = lockMap.remove(key);
        if (removedLock != null && removedLock.isHeldByCurrentThread()) {
            removedLock.unlock();
        }
    }
}
